# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

"""Stackless 状态机测试

测试使用状态机后端的 async 函数生成。
"""

from gtest import EXPECT_EQ, test
from mc.future import ensure_future


# 测试单个 await（状态机后端）
@test("StacklessStateMachine", "SingleAwait")
def test_single_await_state_machine():
    """测试单个 await 的状态机生成"""
    
    async def step() -> int:
        return 7
    
    # @mcpy_async_backend: state_machine
    async def compute() -> int:
        v = await step()
        return v + 5
    
    result = ensure_future(compute())
    EXPECT_EQ(result.get(), 12)


# 测试多个 await（状态机后端）
@test("StacklessStateMachine", "MultipleAwaits")
def test_multiple_awaits_state_machine():
    """测试多个 await 的状态机生成"""
    
    async def s1() -> int:
        return 3
    
    async def s2(x: int) -> int:
        return x * 4
    
    async def s3(x: int) -> int:
        return x + 2
    
    # @mcpy_async_backend: state_machine
    async def pipeline() -> int:
        a = await s1()
        mid = a + 1  # between 语句
        b = await s2(mid)
        tail = b - 2  # between 语句
        c = await s3(tail)
        return c * 2
    
    result = ensure_future(pipeline())
    # 计算期望值：a=3, mid=4, b=16, tail=14, c=16, 返回=32
    EXPECT_EQ(result.get(), 32)

